import operator
import pickle
import matplotlib.pyplot as plt
import os
import re
import numpy as np
import pandas as pd
import sys

VER = 'v5.0'
PICKLE_TAG = 'pickle'
SAVE_DIR = os.path.join('_save', 'vcode_ds_list_files.py', VER)

# load epoch infos
files = os.listdir(SAVE_DIR)
pickle_regexp = re.compile(f'^{PICKLE_TAG}-(\\d+)-([^\\-\\.]+)\\.')
epoch_infos = []
for file_name in files:
    path = os.path.join(SAVE_DIR, file_name)
    base, ext = os.path.splitext(file_name)
    ext = ext.lower()
    matcher = pickle_regexp.match(file_name)
    if matcher is not None:
        info_row = []
        groups = matcher.groups()
        info_row.append(int(groups[0]))
        info_row.append(groups[1])
        with open(path, 'br') as f:
            info = pickle.load(f)
            info_row.append(info)
        epoch_infos.append(info_row)


# sort it
df = pd.DataFrame(epoch_infos)
df = df.sort_values(by=[0, 1], axis=0)

# drill data
keys_loss = ['loss']
keys_acc = ['accuracy', 'sparse_categorical_accuracy', 'my_acc', 'my_acc_digit']
keys_loss[len(keys_loss):] = ['val_' + key for key in keys_loss]
keys_acc[len(keys_acc):] = ['val_' + key for key in keys_acc]
n_epochs = df.iloc[-1, 0]
print(f'n_epochs: {n_epochs}')
tgt_map = dict()
for keys in [keys_loss, keys_acc]:
    for key in keys:
        tgt_map[key] = np.zeros(n_epochs)
for row in df.index:
    info_row = df.iloc[row]
    for keys in [keys_loss, keys_acc]:
        for key in keys:
            tgt_map[key][info_row[0] - 1] = info_row[2][key]

# plot it
spr = 1
spc = 2
spn = 0
plt.figure(figsize=[12, 6])


def plot_keys(keys, title):
    global spn
    spn += 1
    plt.subplot(spr, spc, spn)
    plt.title(title)
    for key in keys:
        plt.plot(tgt_map[key], label=key)
    plt.legend()


plot_keys(keys_loss, 'loss')
plot_keys(keys_acc, 'acc')

# show it
plt.show()
